#include "kernel_operator.h"
#include <type_traits>
#define K_MAX_SHAPE_DIM 0
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2; 


template<typename TYPE_X1, typename TYPE_X2,typename TYPE_Y> class KernelPows {
    public:
        __aicore__ inline KernelPows() {}
        __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, 
            uint32_t ALIGN_NUM, uint32_t block_size, uint32_t core_size, uint32_t core_remain,
            TPipe * pipeIn) {
            ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
            this->blockLength = core_size + (GetBlockNum() == GetBlockIdx() + 1 ? core_remain : 0);
            this->tileLength = block_size;
            this->ALIGN_NUM = ALIGN_NUM;
            this->blockLength = this->blockLength + (this->blockLength % ALIGN_NUM ? ALIGN_NUM - this->blockLength % ALIGN_NUM : 0);
            //this->totalLength = totalLength;
            auto startPointer = core_size * GetBlockIdx();
            auto bufferlength = this->blockLength;

    
            Gm_x1.SetGlobalBuffer((__gm__ TYPE_X1*)x1 + startPointer ,bufferlength);
            Gm_x2.SetGlobalBuffer((__gm__ TYPE_X1*)x2 + startPointer , bufferlength);
            //Gm_cdt.SetGlobalBuffer((__gm__  uint8_t *)condition + startPointer, bufferlength);
            Gm_y.SetGlobalBuffer((__gm__ TYPE_X1*)y + startPointer , bufferlength);
            
    
            this->tileNum = this->blockLength / this->tileLength + (this->blockLength % this->tileLength > 0);
    
            pipe = pipeIn;
            pipe->InitBuffer(Q_x1, BUFFER_NUM, this->tileLength * sizeof(TYPE_X1));
            pipe->InitBuffer(Q_x2, BUFFER_NUM, this->tileLength * sizeof(TYPE_X1));
            //pipe.InitBuffer(Q_cdt, BUFFER_NUM, this->tileLength * sizeof(uint8_t));
            pipe->InitBuffer(Q_y, BUFFER_NUM, this->tileLength * sizeof(TYPE_X1));


            if constexpr (! std::is_same_v<TYPE_X1, float>) {

            pipe->InitBuffer(B_x1, this->tileLength * sizeof(float));
            pipe->InitBuffer(B_x2, this->tileLength * sizeof(float));

            }
            


            
        }
        __aicore__ inline void Process() { //假设每次处理一行来看


           

            

            int32_t loopCount = this->tileNum;
            for (int32_t i = 0; i < loopCount-1; i++) {
                CopyIn(i, this->tileLength);
                Compute(i, this->tileLength);
                CopyOut(i, this->tileLength);
            }
            uint32_t length = this->blockLength - this->tileLength * (loopCount - 1);
            CopyIn(loopCount - 1, (length +31 ) / 32 * 32);
            Compute(loopCount - 1, (length +31 ) / 32 * 32);
            CopyOut(loopCount - 1, (length +31 ) / 32 * 32);
          

       
           
        }   
    
         
    private:
        __aicore__ inline void CopyIn(int32_t progress, uint32_t length) {
           LocalTensor<TYPE_X1> x1 = Q_x1.AllocTensor<TYPE_X1>();
           LocalTensor<TYPE_X1> x2 = Q_x2.AllocTensor<TYPE_X1>();
           //LocalTensor<uint8_t> cdt = Q_cdt.AllocTensor<uint8_t>();
           DataCopy(x1, Gm_x1[progress * this->tileLength], length);
           DataCopy(x2, Gm_x2[progress * this->tileLength], length);
           
           //DataCopy(cdt, (Gm_cdt[progress * this->tileLength]), length);
           
           Q_x1.EnQue(x1);
           Q_x2.EnQue(x2);
           //Q_cdt.EnQue(cdt);
        }
        __aicore__ inline void Compute(int32_t progress, uint32_t length) {
            LocalTensor<TYPE_X1> x1Local = Q_x1.DeQue<TYPE_X1>();
            LocalTensor<TYPE_X1> x2Local = Q_x2.DeQue<TYPE_X1>();
            //LocalTensor<uint8_t> cdt = Q_cdt.DeQue<uint8_t>();
            LocalTensor<TYPE_X1> yLocal = Q_y.AllocTensor<TYPE_X1>();

            //auto bits = B_bits.Get<uint8_t>();

            if constexpr (std::is_same_v<TYPE_X1,float>){
                Ln(x1Local,x1Local,length);
                Mul(x1Local,x1Local,x2Local,length);
                Exp(yLocal,x1Local,length);
    
            }
    
            else if constexpr (std::is_same_v<TYPE_X1,half>){
    
                auto p1 = B_x1.Get<float>();
                auto p2 = B_x2.Get<float>();
    
                Cast(p1, x1Local, RoundMode::CAST_NONE, length);
                Cast(p2, x2Local, RoundMode::CAST_NONE, length);
    
                Ln(p1,p1,length);
                Mul(p1,p1,p2,length);
                Exp(p1,p1,length);
    
                Cast(yLocal,p1,RoundMode::CAST_NONE,length);
    
    
    
    
            }
            else{
    
                auto p1 = B_x1.Get<float>();
                auto p2 = B_x2.Get<float>();
    
                Cast(p1, x1Local, RoundMode::CAST_NONE, length);
                Cast(p2, x2Local, RoundMode::CAST_NONE, length);
    
                Ln(p1,p1,length);
                Mul(p1,p1,p2,length);
                Exp(p1,p1,length);
    
                Cast(yLocal,p1,RoundMode::CAST_RINT,length);
    
    
    
            }

            
            Q_x1.FreeTensor(x1Local);  
            Q_x2.FreeTensor(x2Local);
            //Q_cdt.FreeTensor(cdt);
            Q_y.EnQue<TYPE_X1>(yLocal);
        }
        __aicore__ inline void CopyOut(int32_t progress, uint32_t length) {
            LocalTensor<TYPE_X1> y = Q_y.DeQue<TYPE_X1>();
            DataCopy(Gm_y[progress * this->tileLength], y, length);
            Q_y.FreeTensor(y);
        }
    private:
        TPipe * pipe;
        TQue<QuePosition::VECIN, BUFFER_NUM> Q_x1, Q_x2;
        TQue<QuePosition::VECOUT, BUFFER_NUM> Q_y;

        //TBuf<QuePosition::VECCALC> B_result, B_zero, B_bits,B_half;
        //TBuf<QuePosition::VECCALC> B_x1, B_x2;
        TBuf<QuePosition::VECCALC> B_x1, B_x2;

        // int32_t M;
        // int32_t N;
        // int32_t Z;

        // LocalTensor<half> zero;
      
        GlobalTensor<TYPE_X1> Gm_x1, Gm_x2, Gm_y;
        //GlobalTensor<uint8_t> Gm_cdt;
        uint32_t blockLength;
        uint32_t tileNum;
        //uint32_t totalLength;
        uint32_t tileLength;
        uint32_t ALIGN_NUM;
    };

template<typename TYPE_X1, typename TYPE_X2,typename TYPE_Y> class Pows_Broadcast {

    using T = TYPE_X1;
    public:
        __aicore__ inline Pows_Broadcast() {}
        __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
            uint32_t ALIGN_NUM, uint32_t block_size, uint32_t core_size, uint32_t core_remain,
            uint32_t total_length, uint32_t x1_length, uint32_t x2_length, 
            int32_t ss[], int32_t numshapes, int32_t sf[]) {

                ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

                // this->coreDataNum = CoreDataNum;
                // this->tileNum = finalTileNum;
                // this->tileDataNum = tileDataNum;
                // this->tailDataNum = TailDataNum;

                this->blockLength = core_size + (GetBlockNum() == GetBlockIdx() + 1 ? core_remain : 0);
                this->tileLength = block_size;
                this->ALIGN_NUM = ALIGN_NUM;
                this->blockLength = this->blockLength + (this->blockLength % ALIGN_NUM ? ALIGN_NUM - this->blockLength % ALIGN_NUM : 0);
                //this->totalLength = totalLength;
                auto startPointer = core_size * GetBlockIdx();
                auto bufferlength = this->blockLength;
                for (int i = 0; i < 128; ++i) {
                    ((int32_t *)this->shape)[i] = ss[i];
                }
                this->numshapes = numshapes;
                for(int i = 0; i < 64; ++i) {
                    ((int32_t *)this->shapefull)[i] = sf[i];
                }
                this->total_length = total_length;
                this->x1_Length = x1_length;
                this->x2_Length = x2_length;

                
        
                x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x1+startPointer, bufferlength);
                x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2*)x2+startPointer, bufferlength);
                yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y+startPointer, bufferlength);

                this->tileNum = this->blockLength / this->tileLength + (this->blockLength % this->tileLength > 0);
        
                pipe.InitBuffer(inQueueX1, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X1));
                pipe.InitBuffer(inQueueX2, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X2));
                pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(DTYPE_Y));

                if constexpr (! std::is_same_v<TYPE_X1, float>) {
                     pipe.InitBuffer(QueueTmp1, this->tileLength * sizeof(float));
                     pipe.InitBuffer(QueueTmp2, this->tileLength * sizeof(float));
                }
        }
        __aicore__ inline void Process() {

            int32_t loopCount = this->tileNum;
            for (int32_t i = 0; i < loopCount-1; i++) {
                CopyIn(i, this->tileLength);
                Compute(i, this->tileLength);
                CopyOut(i, this->tileLength);
            }
            uint32_t length = this->blockLength - this->tileLength * (loopCount - 1);
            CopyIn(loopCount - 1, (length +31 ) / 32 * 32);
            Compute(loopCount - 1, (length +31 ) / 32 * 32);
            CopyOut(loopCount - 1, (length +31 ) / 32 * 32);
        }
    
    private:
    __aicore__ inline void CopyIn(int32_t progress,uint32_t length)
    {
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
        LocalTensor<DTYPE_X2> x2Local = inQueueX2.AllocTensor<DTYPE_X2>();
        
        if(this->x1_Length < this->total_length) {
            BroadCX1(x1Local, progress * this->tileLength, length);
        } else {            
            DataCopy(x1Local, x1Gm[progress * this->tileLength], length);
        }
        if(this->x2_Length < this->total_length) {
            BroadCX2(x2Local, progress * this->tileLength, length);
        } else {
            DataCopy(x2Local, x2Gm[progress * this->tileLength], length);
        }
        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
    }

    __aicore__ inline void BroadCX1(LocalTensor<DTYPE_X1> &dst, uint32_t offset, uint32_t length) {
        if(this->x1_Length == 1) {

            DTYPE_X1 tmp = x1Gm.GetValue(0);

            for(int i=0;i<length;i++){

                dst.SetValue(i,tmp);

            }
            
            //Duplicate(dst, tmp, length);
            return;
        }
        for(uint32_t i = 0; i < length; i++) {
            int idxtmp = 0;
            int istart = i + offset;
            for(int k = 1; k <= this->numshapes; k++) {
                int kpos = 0;
                int krange = 1;
                if(k < this->numshapes) {
                    for(int m = k + 1; m <= this->numshapes; m++) {
                        krange *= shapefull[m - 1];
                    }
                    kpos = istart / krange;
                    istart = istart % krange;
                } else {
                    krange = shapefull[k - 1];
                    kpos = istart % krange;
                }
                //idxtmp += kpos * this->stride[k - 1];
                int krangeB = 1;
                if(shapefull[k - 1] == shape[0][k - 1]) {
                    if(k < this->numshapes) {
                        for(int m = k + 1; m <= this->numshapes; m++) {
                            krangeB *= shape[0][m - 1];
                        }
                        idxtmp += kpos * krangeB;
                    }  else {
                        idxtmp += kpos;
                    }
                }
            }
            DTYPE_X1 tmp = x1Gm.GetValue(idxtmp);
            dst.SetValue(i, tmp);
        }
    }
    __aicore__ inline void BroadCX2(LocalTensor<DTYPE_X1> &dst, uint32_t offset, uint32_t length) {
        if(this->x2_Length == 1) {
            DTYPE_X1 tmp = x2Gm.GetValue(0);
            //Duplicate(dst, tmp, length);

            for(int i=0;i<length;i++){

                dst.SetValue(i,tmp);

            }

            return;
        }
        for(uint32_t i = 0; i < length; i++) {
            int idxtmp = 0;
            int istart = i + offset;
            for(int k = 1; k <= this->numshapes; k++) {
                int kpos = 0;
                int krange = 1;
                if(k < this->numshapes) {
                    for(int m = k + 1; m <= this->numshapes; m++) {
                        krange *= shapefull[m - 1];
                    }
                    kpos = istart / krange;
                    istart = istart % krange;
                } else {
                    krange = shapefull[k - 1];
                    kpos = istart % krange;
                }
                //idxtmp += kpos * this->stride[k - 1];
                int krangeB = 1;
                if(shapefull[k - 1] == shape[1][k - 1]) {
                    if(k < this->numshapes) {
                        for(int m = k + 1; m <= this->numshapes; m++) {
                            krangeB *= shape[1][m - 1];
                        }
                        idxtmp += kpos * krangeB;
                    }  else {
                        idxtmp += kpos;
                    }
                }
            }
            DTYPE_X1 tmp = x2Gm.GetValue(idxtmp);
            dst.SetValue(i, tmp);
        }
    }
        __aicore__ inline void Compute(uint32_t progress,uint32_t length) {
            LocalTensor<DTYPE_X1> x1Local = inQueueX1.DeQue<DTYPE_X1>();
            LocalTensor<DTYPE_X2> x2Local = inQueueX2.DeQue<DTYPE_X2>();
            LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
           
            // Ln(x1LocalB,x1LocalB,length);
            // Mul(x1LocalB,x1LocalB,x2LocalB,length);
            // Exp(yLocalB,x1LocalB,length);
        
            if constexpr (std::is_same_v<TYPE_X1,float>){
                Ln(x1Local,x1Local,length);
                Mul(x1Local,x1Local,x2Local,length);
                Exp(yLocal,x1Local,length);
    
            }
    
            else if constexpr (std::is_same_v<TYPE_X1,half>){

                auto p1 = QueueTmp1.Get<float>();
                auto p2 = QueueTmp2.Get<float>();
    
                //  p1 = calbuf1.Get<float>();
                //  p2 = calbuf2.Get<float>();
    
                Cast(p1, x1Local, RoundMode::CAST_NONE, length);
                Cast(p2, x2Local, RoundMode::CAST_NONE, length);
    
                Ln(p1,p1,length);
                Mul(p1,p1,p2,length);
                Exp(p1,p1,length);
    
                Cast(yLocal,p1,RoundMode::CAST_NONE,length);
    
    
    
    
            }
            else{
    
                auto p1 = QueueTmp1.Get<float>();
                auto p2 = QueueTmp2.Get<float>();
    
                Cast(p1, x1Local, RoundMode::CAST_NONE, length);
                Cast(p2, x2Local, RoundMode::CAST_NONE, length);
    
                Ln(p1,p1,length);
                Mul(p1,p1,p2,length);
                Exp(p1,p1,length);
    
                Cast(yLocal,p1,RoundMode::CAST_RINT,length);
    
    
    
            }
            
            outQueueY.EnQue<DTYPE_Y>(yLocal);
            inQueueX1.FreeTensor(x1Local);
            inQueueX2.FreeTensor(x2Local);
            //inQueueWEIGHTB.FreeTensor(weightLocalB);
        }
       
        __aicore__ inline void CopyOut(int32_t progress,uint32_t length)
    {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        DataCopy(yGm[progress * this->tileLength], yLocal, length);
        outQueueY.FreeTensor(yLocal);
    }
        
    private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX1, inQueueX2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2;

    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_X2> x2Gm;
    GlobalTensor<DTYPE_Y> yGm;

    uint32_t blockLength;
    uint32_t tileNum;
    //uint32_t totalLength;
    uint32_t tileLength;
    uint32_t ALIGN_NUM;
    
    uint32_t x1_Length;
    uint32_t x2_Length;
    uint32_t total_length;
    int32_t shape[2][64];
    int32_t numshapes;
    int32_t shapefull[64];
    };



extern "C" __global__ __aicore__ void pows(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);


    if(TILING_KEY_IS(1)){
        TPipe pipe;
        KernelPows<DTYPE_X1, DTYPE_X2, DTYPE_Y> op;
        op.Init(x1, x2, y,
            tiling_data.ALIGN_NUM, tiling_data.block_size, tiling_data.core_size, tiling_data.core_remain,&pipe);
        op.Process();

    }
    else if(TILING_KEY_IS(2)){

        Pows_Broadcast<DTYPE_X1, DTYPE_X2, DTYPE_Y> op;
        op.Init(x1, x2, y,
            tiling_data.ALIGN_NUM, tiling_data.block_size, tiling_data.core_size, tiling_data.core_remain,
            tiling_data.total_length, tiling_data.x1_length, tiling_data.x2_length, 
            tiling_data.shape, tiling_data.numshapes, tiling_data.shapefull
        );
        op.Process();



    }


    // TODO: user kernel impl
}